from spk.core.workspace import *
import streamlit as st

import requests
from urllib.parse import parse_qs, urlsplit
import json
from datetime import datetime

def get_url_param(url: str, param: str) -> str | None:
    return parse_qs(urlsplit(url).query).get(param, [None])[0]

def get_date(article):
    date = datetime.fromtimestamp(int(article['create_time'])).strftime('%Y-%m-%d')
    return date

class StepWeixinLinkToDataset(SpkWorker):
    '''Weixin 文章列表链接转换

    这是一个简单的工作流
    1. 输入微信链接
    2. 解析微信链接
    3. 制作微信数据集
    4. 显示微信数据集
    5. 下载微信数据集
    '''

    First  = '第一步，输入微信链接'
    Second = '第二步，解析微信链接'
    Three  = '第三步，制作微信数据集'
    Four   = '第四步，显示微信数据集'
    Five   = '第五步，下载微信数据集'

    __steps__ = [
        First, Second, Three, Four, Five
    ]

    __keywords__ = [
        'spk-weixin-link-input',
        'spk-weixin-link-request',
        'spk-weixin-link-dataset',
    ]

    def __do_work__(self):

        # 第一步，输入微信链接
        if current() == self.First:
            weixin_link = st.text_input('微信链接', '')
            if st.button('解析微信链接'):
                head = requests.head(weixin_link)
                if head.status_code == 200:
                    st.success('微信链接有效')
                    set('spk-weixin-link-input', weixin_link)
                    next_step()


            # 横向排列
            with st.expander('微信链接示例', expanded=True):
                    # 创建两列布局
                col1, col2 = st.columns([1,2])
                
                with col1:
                    if st.button('显示 bash 脚本', help="点击查看采集脚本的Bash实现"):
                        self.bash_script_dialog()
                
                with col2:
                    if st.button('显示 python 脚本', help="点击查看采集脚本的Python实现"):
                        self.python_script_dialog()

        # 第二步，解析微信链接
        if current() == self.Second:
            
            if st.button('上一步'):
                set('spk-weixin-link-input', None)
                prev_step()

            if get('spk-weixin-link-input'):
                st.write('微信链接：')
                st.write(get('spk-weixin-link-input'))


            biz = get_url_param(get('spk-weixin-link-input'), '__biz')
            album_id = get_url_param(get('spk-weixin-link-input'), 'album_id')

            st.table({
                '参数': ['__biz', 'album_id'],
                '值': [biz, album_id]
            })

            if st.button('制作数据集'):
                next_step()


        # 第三步，制作微信数据集
        if current() == self.Three:
            if st.button('上一步'):
                set('spk-weixin-link-request', None)
                prev_step()

            if get('spk-weixin-link-request') is None:

                # 初始化数据集合
                all_articles = []
                session = requests.Session()

                biz = get_url_param(get('spk-weixin-link-input'), '__biz')
                album_id = get_url_param(get('spk-weixin-link-input'), 'album_id')
                base_url = "https://mp.weixin.qq.com/mp/appmsgalbum"
                def make_request(params):
                    try:
                        resp = session.get(
                            base_url,
                            params={**params, "__biz": biz, "album_id": album_id, "f": "json"},
                            timeout=10
                        )
                        resp.raise_for_status()
                        return resp.json()
                    except Exception as e:
                        print(f"Request failed: {e}")
                        return None
                

                with st.spinner('正在请求初始数据...'):
                    # 初始请求
                    initial_data = make_request({"action": "getalbum"})
                    if not initial_data:
                        st.error('请求初始数据失败')
                        return
                    st.success('请求初始数据成功')
                
                # 提取第一页数据
                articles = initial_data.get("getalbum_resp", {}).get("article_list", [])
                all_articles.extend(articles)
                print(f"Initial articles fetched: {len(articles)}")

                # 获取分页参数
                last_article = articles[-1] if articles else None
                msgid = last_article.get("msgid") if last_article else None
                itemidx = last_article.get("itemidx") if last_article else None
                
                request_count = 0
                # 分页请求
                with st.spinner('正在请求分页数据，请稍等...'):
                    while msgid and msgid != "null":
                        params = {
                            "action": "getalbum",
                            "count": 10,
                            "begin_msgid": msgid,
                            "begin_itemidx": itemidx
                        }
                        
                        page_data = make_request(params)
                        if not page_data:
                            break

                        new_articles = page_data.get("getalbum_resp", {}).get("article_list", [])
                        if not new_articles:
                            break
                        if type(new_articles) != list:
                            new_articles = [new_articles]

                        all_articles.extend(new_articles)
                        print(f"Fetched {len(new_articles)} new articles, total: {len(all_articles)}")

                        # 更新分页锚点
                        last_article = new_articles[-1]
                        msgid = last_article.get("msgid")
                        itemidx = last_article.get("itemidx")
                        request_count += 1
                st.write('已发出请求次数：' + str(request_count))
                st.success('请求分页数据成功')
                
                st.write('总文章数：' + str(len(all_articles)))
                set('spk-weixin-link-dataset', all_articles)
                set('spk-weixin-link-request', True)

            if get('spk-weixin-link-dataset') is not None:
                if st.button('显示数据集'):
                    next_step()

        # 第四步，显示数据集
        if current() == self.Four:

            if st.button('上一步'):
                prev_step()

            st.write('总文章数：' + str(len(get('spk-weixin-link-dataset'))))

            tabs = st.tabs(['Table', 'JSON', 'Code', 'Title/Link/Code', 'Wiki-数据集'])
            with tabs[0]:
                st.table({
                    '标题': [article['title'] for article in get('spk-weixin-link-dataset')],
                })
            with tabs[1]:
                st.json(get('spk-weixin-link-dataset'))
            with tabs[2]:
                st.code(json.dumps(get('spk-weixin-link-dataset'), indent=4, ensure_ascii=False))
            with tabs[3]:
                objs = [{'title': article['title'], 'url': article['url']} for article in get('spk-weixin-link-dataset')]
                st.code(json.dumps(objs, indent=4, ensure_ascii=False))
            with tabs[4]:
                wiki_tabs = st.tabs(['TextArea', 'TextArea/Date', 'Codes', 'CodeJson'])
                with wiki_tabs[0]:
                    wiki_columns = st.columns(2)
                    with wiki_columns[0]:
                        st.write('标题')
                        st.text_area(label="wiki 标题/日期", value='\n'.join([article['title'] for article in get('spk-weixin-link-dataset')]), height=400)
                    with wiki_columns[1]:
                        st.write('链接')
                        st.text_area(label="wiki 链接/日斯", value='\n'.join([article['url'] for article in get('spk-weixin-link-dataset')]), height=400)
                with wiki_tabs[1]:
                    wiki_columns = st.columns(2)
                    with wiki_columns[0]:
                        st.write('标题')
                        st.text_area(label="wiki 标题", value='\n'.join([f"[{get_date(article)}] {article['title']}" for article in get('spk-weixin-link-dataset')]), height=400)
                    with wiki_columns[1]:
                        st.write('链接')
                        st.text_area(label="wiki 链接", value='\n'.join([article['url'] for article in get('spk-weixin-link-dataset')]), height=400)
                with wiki_tabs[2]:
                    wiki_columns = st.columns(2)
                    with wiki_columns[0]:
                        st.write('标题')
                        st.code('\n'.join([article['title'] for article in get('spk-weixin-link-dataset')]))
                    with wiki_columns[1]:
                        st.write('链接')
                        st.code('\n'.join([article['url'] for article in get('spk-weixin-link-dataset')]))
                with wiki_tabs[3]:
                    wiki_columns = st.columns(2)
                    with wiki_columns[0]:
                        st.write('标题')
                        st.code(json.dumps([article['title'] for article in get('spk-weixin-link-dataset')], indent=4, ensure_ascii=False))
                    with wiki_columns[1]:
                        st.write('链接')
                        st.code(json.dumps([article['url'] for article in get('spk-weixin-link-dataset')], indent=4, ensure_ascii=False))

            if st.button('下载数据集'):
                next_step()

        # 第五步，下载数据集
        if current() == self.Five:
            if st.button('上一步'):
                prev_step()
        
            if get('spk-weixin-link-dataset'):
                st.write('总文章数：' + str(len(get('spk-weixin-link-dataset'))))

                st.download_button(
                    label="下载数据集",
                    data=json.dumps(get('spk-weixin-link-dataset'), indent=4, ensure_ascii=False),
                    file_name="weixin_articles.json",
                    mime="application/json"
                )



    @st.dialog("可用 shell 脚本", width="large")
    def bash_script_dialog(self):
        with open('spk/plugins/tests/weixin_test.sh', 'r') as f:
            st.code(f.read(), language='bash')

    @st.dialog("可用 python 脚本", width="large")
    def python_script_dialog(self):
        with open('spk/plugins/tests/weixin_test.py', 'r') as f:
            st.code(f.read(), language='python')